import os
import math
from PIL import Image

class VectorCompare:
    # 计算矢量大小
    def magnitude(self, concordance):
        total = 0
        for word, count in concordance.items():
            total += count ** 2
        return math.sqrt(total)
    
    # 计算矢量之间的 cos 值
    def relation(self, concordance1, concordance2):
        relevance = 0
        topvalue = 0
        for word, count in concordance1.items():
            if word in concordance2:
                topvalue += count * concordance2[word]
        return topvalue / (self.magnitude(concordance1) * self.magnitude(concordance2))

# 将图片转换为矢量
def buildvector(img):
    d1 = {}
    count = 0
    for i in img.getdata():
        d1[count] = i
        count += 1
    return d1

v = VectorCompare()
iconset = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f', 'g', \
    'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']

# 加载数据集
imageset = []
for letter in iconset:
    for img in os.listdir('iconset/{}'.format(letter)):
        tmp = []
        tmp.append(buildvector(Image.open('iconset/{}/{}'.format(letter, img))))
        imageset.append({letter: tmp})

img = Image.open('captcha.gif')
# （将图片转换为 8 位像素模式）
img.convert('P')
img2 = Image.new('P', img.size, 255)

for x in range(img.size[1]):
    for y in range(img.size[0]):
        pix = img.getpixel((y, x))
        if pix == 220 or pix == 227:  # these are the numbers to get
            img2.putpixel((y, x), 0)
img2.save('new.gif')

inletter = False
foundletter = False
start = 0
end = 0

letters = []

for y in range(img2.size[0]):
    for x in range(img2.size[1]):
        pix = img2.getpixel((y, x))
        if pix != 255:
            inletter = True
    
    if foundletter == False and inletter == True:
        foundletter = True
        start = y
    
    if foundletter == True and inletter == False:
        foundletter = False
        end = y
        letters.append((start, end))
    
    inletter = False

count = 0
for letter in letters:
    # 前两个值为左上角坐标
    # 后两个值为右下角坐标
    img3 = img2.crop((letter[0], 0, letter[1], img2.size[1]))
    img3.save('{}.gif'.format(count))

    guess = []
    for image in imageset:
        for x, y in image.items():
            if len(y) != 0:
                guess.append((v.relation(y[0], buildvector(img3)), x))
    guess.sort(reverse=True)
    print(guess[0])

    count += 1
